import os
from functools import cached_property
import logging
from typing import Any
from typing_extensions import Self
from pydantic import BaseModel, Field
from pysam import FastaFile, VariantFile, VariantRecord, TabixFile  # pylint: disable=E0611
from .gene import Region, RegionType, Transcript
from .seq import Sequence, MTSequence

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("OpenAnno SNV")


class TranscriptAnnotation(BaseModel):
    gene: str
    gene_id: str
    transcript: str
    region: str = Field(default="")
    na_change: str = Field(default="")
    aa_change: str = Field(default="")
    event: str = Field(default="")
    region2: str = Field(default="")

    @cached_property
    def detail(self) -> str:
        text = ""
        if self.na_change:
            text = f"{self.gene}:{self.transcript}:{self.region2}:{self.na_change}"
            if self.aa_change:
                text += f":{self.aa_change}"
        return text

    @classmethod
    def from_transcript(cls, transcript: Transcript, *regions: Region) -> Self:
        instance = cls(gene=transcript.gene, gene_id=transcript.gene_id, transcript=transcript.name)
        regions = sorted(regions, key=lambda x: x.start if transcript.strand == "+" else -x.start)
        if regions:
            region1, region2 = regions[0], regions[-1]
            instance.region2 = region1.exon if region1.exon == region2.exon else f"{region1.exon}_{region2.exon}"
            region_types = list(map(lambda x: x.type, regions))
            if RegionType.CDS in region_types:
                instance.region = "exonic"
            elif RegionType.UTR in region_types:
                current = regions[region_types.index(RegionType.UTR)]
                instance.region = current.name
            elif RegionType.INTRON in region_types:
                instance.region = "intronic"
            elif RegionType.UD in region_types:
                current = regions[region_types.index(RegionType.UD)]
                instance.region = (
                    "upstream"
                    if (current.end < transcript.tx_start and transcript.strand == "+") or (current.start > transcript.tx_end and transcript.strand == "-")
                    else "downstream"
                )
        return instance


class SNV(BaseModel):
    chrom: str
    pos: int
    start: int
    end: int
    ref: Sequence | MTSequence
    alt: Sequence | MTSequence
    info: dict[str, Any] = Field(default_factory=dict)

    class Config:
        arbitrary_types_allowed = True

    @cached_property
    def pk(self) -> str:
        return f"{self.chrom}:{self.pos}:{self.ref}:{self.alt}"

    @cached_property
    def s_start(self) -> int:
        """standard start"""
        raise NotImplementedError("This method should be implemented.")

    @cached_property
    def s_end(self) -> int:
        """standard end"""
        return self.end

    @cached_property
    def s_ref(self) -> Sequence | MTSequence:
        """standard ref"""
        raise NotImplementedError("This method should be implemented.")

    @cached_property
    def s_alt(self) -> Sequence | MTSequence:
        """standard alt"""
        raise NotImplementedError("This method should be implemented.")

    @classmethod
    def from_variant_record(cls, record: VariantRecord) -> Self:
        seq_cls = Sequence if record.chrom in ["chrM", "MT"] else MTSequence
        return cls(
            chrom=record.chrom,
            pos=record.pos,
            start=record.start,
            end=record.stop,
            ref=seq_cls(record.ref),
            alt=seq_cls(record.alts[0]),
        )

    def annotate_mrna(self, transcript: Transcript) -> TranscriptAnnotation:
        raise NotImplementedError("This method should be implemented.")

    def annotate_ncrna(self, transcript: Transcript) -> TranscriptAnnotation:
        raise NotImplementedError("This method should be implemented.")

    def annotate_gene_based(self, genepred_tabix: TabixFile, genome_fasta: FastaFile, gene_id_dict: dict[str, str]):
        try:
            lines = genepred_tabix.fetch(self.chrom, self.start, self.end)
        except:
            logger.warning("'%s' not found in %s", self.chrom, genepred_tabix.filename.decode())
        else:
            trans_anno_dict: dict[str, list[TranscriptAnnotation]] = {}
            for line in lines:
                transcript = Transcript.from_genepred(line)
                if transcript.tx_start <= self.s_start and transcript.tx_end >= self.s_end:
                    transcript.with_data(genome_fasta, gene_id_dict)
                    trans_anno = self.annotate_ncrna(transcript) if transcript.is_unk else self.annotate_mrna(transcript)
                    trans_anno_dict.setdefault(trans_anno.gene, []).append(trans_anno)
            genes, gene_ids, events, regions, details = [], [], [], [], []
            for gene, trans_anno_list in trans_anno_dict.items():
                gene_id = ""
                _events, _regions1, _regions2, _details = set(), set(), set(), set()
                for trans_anno in trans_anno_list:
                    if not gene_id and trans_anno.gene_id:
                        gene_id = trans_anno.gene_id
                    if trans_anno.event:
                        _events.add(trans_anno.event)
                    if trans_anno.detail:
                        _details.add(trans_anno.detail)
                    if trans_anno.region in ["exonic", "splicing", "exonic_splicing", "transcript"]:
                        _regions1.add(trans_anno.region)
                    elif trans_anno.region in ["ncRNA", "UTR3", "UTR5", "intronic"]:
                        _regions2.add(trans_anno.region)
                genes.append(gene)
                gene_ids.append(gene_id or ".")
                events.append("|".join(_events) or ".")
                details.append("|".join(_details) or ".")
                regions.append("|".join(_regions1) or "|".join(_regions2) or ".")
            self.info.update(
                {
                    "GENE": genes or ["."],
                    "GENE_ID": gene_ids or ["."],
                    "EVENT": events or ["."],
                    "REGION": regions or ["."],
                    "DETAIL": details or ["."],
                }
            )

    def annotate_filter_based(self, variant_file: VariantFile):
        try:
            records = variant_file.fetch(self.chrom, self.start, self.end)
        except:
            logger.warning("'%s' not found in %s", self.chrom, variant_file.filename.decode())
        else:
            for record in filter(lambda x: x.pos == self.pos and x.ref == self.ref and x.alts and x.alts[0] == self.alt, records):
                self.info.update(dict(record.info))

    def annotate_region_based(self, tabix_file: TabixFile):
        try:
            lines = tabix_file.fetch(self.chrom, self.start, self.end)
        except:
            logger.warning("'%s' not found in %s", self.chrom, tabix_file.filename.decode())
        else:
            val = "|".join(list(map(lambda x: x.split("\t")[3], lines)))
            if val:
                name = os.path.basename(tabix_file.filename.decode()).split(".")[0]
                self.info[name] = val
